// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <stdint.h>
#include <string.h>

#include <melon/executors/cpu_thread_pool_executor.h>

#include <pollux/common/compression/compression.h>

namespace kumo::pollux::common {
    /// Provides the fine-grained spill execution stats.
    struct SpillStats {
        /// The number of times that spilling runs on an operator.
        uint64_t spillRuns{0};
        /// The number of bytes in memory to spill
        uint64_t spilledInputBytes{0};
        /// The number of bytes spilled to disks.
        ///
        /// NOTE: if compression is enabled, this counts the compressed bytes.
        uint64_t spilledBytes{0};
        /// The number of spilled rows.
        uint64_t spilledRows{0};
        /// NOTE: when we sum up the stats from a group of spill operators, it is
        /// the total number of spilled partitions X number of operators.
        uint32_t spilledPartitions{0};
        /// The number of spilled files.
        uint64_t spilledFiles{0};
        /// The time spent on filling rows for spilling.
        uint64_t spillFillTimeNanos{0};
        /// The time spent on sorting rows for spilling.
        uint64_t spillSortTimeNanos{0};
        /// The time spent on extracting vector from RowContainer for spilling.
        uint64_t spillExtractVectorTimeNanos{0};
        /// The time spent on serializing rows for spilling.
        uint64_t spillSerializationTimeNanos{0};
        /// The number of spill writer flushes, equivalent to number of write calls to
        /// underlying filesystem.
        uint64_t spillWrites{0};
        /// The time spent on copy out serialized rows for disk write. If compression
        /// is enabled, this includes the compression time.
        uint64_t spillFlushTimeNanos{0};
        /// The time spent on writing spilled rows to disk.
        uint64_t spillWriteTimeNanos{0};
        /// The number of times that an hash build operator exceeds the max spill
        /// limit.
        uint64_t spillMaxLevelExceededCount{0};
        /// The number of bytes read from spilled files.
        uint64_t spillReadBytes{0};
        /// The number of spill reader reads, equivalent to the number of read calls
        /// to the underlying filesystem.
        uint64_t spillReads{0};
        /// The time spent on read data from spilled files.
        uint64_t spillReadTimeNanos{0};
        /// The time spent on deserializing rows read from spilled files.
        uint64_t spillDeserializationTimeNanos{0};

        SpillStats(
            uint64_t _spillRuns,
            uint64_t _spilledInputBytes,
            uint64_t _spilledBytes,
            uint64_t _spilledRows,
            uint32_t _spilledPartitions,
            uint64_t _spilledFiles,
            uint64_t _spillFillTimeNanos,
            uint64_t _spillSortTimeNanos,
            uint64_t _spillExtractVectorTimeNanos,
            uint64_t _spillSerializationTimeNanos,
            uint64_t _spillWrites,
            uint64_t _spillFlushTimeNanos,
            uint64_t _spillWriteTimeNanos,
            uint64_t _spillMaxLevelExceededCount,
            uint64_t _spillReadBytes,
            uint64_t _spillReads,
            uint64_t _spillReadTimeNanos,
            uint64_t _spillDeserializationTimeNanos);

        SpillStats() = default;

        bool empty() const {
            return spilledBytes == 0;
        }

        SpillStats &operator+=(const SpillStats &other);

        SpillStats operator-(const SpillStats &other) const;

        bool operator==(const SpillStats &other) const;

        bool operator!=(const SpillStats &other) const {
            return !(*this == other);
        }

        bool operator>(const SpillStats &other) const;

        bool operator<(const SpillStats &other) const;

        bool operator>=(const SpillStats &other) const;

        bool operator<=(const SpillStats &other) const;

        void reset();

        std::string toString() const;
    };

    MELON_ALWAYS_INLINE std::ostream &operator<<(
        std::ostream &o,
        const common::SpillStats &stats) {
        return o << stats.toString();
    }

    /// The utilities to update the process wide spilling stats.
    /// Updates the number of spill runs.
    void updateGlobalSpillRunStats(uint64_t numRuns);

    /// Updates the stats of new append spilled rows including the number of spilled
    /// rows and the serializaion time.
    void updateGlobalSpillAppendStats(
        uint64_t numRows,
        uint64_t serializaionTimeNs);

    /// Increments the number of spilled partitions.
    void incrementGlobalSpilledPartitionStats();

    /// Updates the time spent on filling rows to spill.
    void updateGlobalSpillFillTime(uint64_t timeNs);

    /// Updates the time spent on sorting rows to spill.
    void updateGlobalSpillSortTime(uint64_t timeNs);

    /// Updates the time spent on extracting vector from RowContainer to spill.
    void updateGlobalSpillExtractVectorTime(uint64_t timeNs);

    /// Updates the stats for disk write including the number of disk writes,
    /// the written bytes, the time spent on copying out (compression) for disk
    /// writes, the time spent on disk writes.
    void updateGlobalSpillWriteStats(
        uint64_t spilledBytes,
        uint64_t flushTimeNs,
        uint64_t writeTimeNs);

    /// Updates the stats for disk read including the number of disk reads, the
    /// amount of data read in bytes, and the time it takes to read from the disk.
    void updateGlobalSpillReadStats(
        uint64_t spillReads,
        uint64_t spillReadBytes,
        uint64_t spillRadTimeNs);

    /// Increments the spill memory bytes.
    void updateGlobalSpillMemoryBytes(uint64_t spilledInputBytes);

    /// Increments the spilled files by one.
    void incrementGlobalSpilledFiles();

    /// Increments the exceeded max spill level count.
    void updateGlobalMaxSpillLevelExceededCount(
        uint64_t maxSpillLevelExceededCount);

    /// Increments the spill read deserialization time.
    void updateGlobalSpillDeserializationTimeNs(uint64_t timeNs);

    /// Gets the cumulative global spill stats.
    SpillStats globalSpillStats();
} // namespace kumo::pollux::common

template<>
struct fmt::formatter<kumo::pollux::common::SpillStats>
        : fmt::formatter<std::string> {
    auto format(const kumo::pollux::common::SpillStats &s, format_context &ctx)
    const {
        return formatter<std::string>::format(s.toString(), ctx);
    }
};
